import os
from typing import Dict, List, Tuple

import pandas as pd

from iflytech_assistant.spark_client import instruct

from .dataclasses import FiveSuggestions, PreferenceTag, RefineWord
from .markdown_prompt_parser import PromptTemplate

__all__ = ["generate", "generate_response"]

templates_dir = os.path.join(os.path.dirname(__file__), "templates")
template_file = os.path.join(templates_dir, "romance_advisor.md")
prompt_template = PromptTemplate.from_markdown_text(open(template_file).read())

SYSTEM_PROMPT = prompt_template.system_prompt_template.raw_str
REPLY_USER_PROMPT = prompt_template.reply_user_prompt_template.raw_str
POLISH_USER_PROMPT = prompt_template.polish_user_prompt_template.raw_str


def build_prompt(
    user_input: str,
    target: str,
    tags: List[PreferenceTag],
    mode: str,
    refine_words: RefineWord = None,
) -> Tuple[str, str]:
    tone_prompt = "\n".join([f"- {tag.tag}: {tag.tone}" for tag in tags if tag.tone])
    style_prompt = "\n".join([f"- {tag.tag}: {tag.style}" for tag in tags if tag.style])
    personality_prompt = "\n".join(
        [f"- {tag.tag}: {tag.personality}" for tag in tags if tag.personality]
    )

    if refine_words:
        if refine_words.type == "tone":
            tone_prompt = refine_words.word
        if refine_words.type == "style":
            style_prompt = refine_words.word

    system_prompt = SYSTEM_PROMPT.format(
        tone=tone_prompt, style=style_prompt, personality=personality_prompt
    )

    if mode == "reply":
        user_prompt = REPLY_USER_PROMPT.format(
            target=target,
            user_input=user_input,
            preferenced_style="、".join([tag.tag for tag in tags]),
        )
    else:
        user_prompt = POLISH_USER_PROMPT.format(
            target=target,
            user_input=user_input,
            preferenced_style="、".join([tag.tag for tag in tags]),
        )
    return system_prompt, user_prompt


def generate(
    user_input,
    target: str,
    tags: List[PreferenceTag],
    mode: str,
    refine_words: RefineWord = None,
):
    system_prompt, user_prompt = build_prompt(
        user_input, target, tags, mode, refine_words
    )
    llm_response = instruct(user_prompt, system_prompt, prefix="", streaming=False)

    suggestions = FiveSuggestions.from_llm_response(llm_response)

    return system_prompt, user_prompt, llm_response, suggestions


def generate_response(
    user_input: str,
    target: str,
    tags: List[PreferenceTag],
    refine_words: RefineWord,
    mode: str,
) -> Tuple[FiveSuggestions, FiveSuggestions]:
    system_prompt, user_prompt, _, suggestions = generate(
        user_input, target, tags, mode
    )

    refine_prompt, _, _, refine_suggestions = generate(
        user_input, target, tags, mode, refine_words
    )
    return (
        suggestions,
        refine_suggestions,
        system_prompt,
        user_prompt,
        refine_prompt,
    )


def load_tags_from_excel(
    file_path: str,
) -> Tuple[Dict[str, List[PreferenceTag]], Dict[str, List[RefineWord]]]:
    # read all sheet names
    xls = pd.ExcelFile(file_path)
    sheets = xls.sheet_names

    assert "tags" in sheets, "tags sheet not found"
    assert "refine_words" in sheets, "refine_words sheet not found"

    refine_words_sheet = xls.parse("refine_words")
    # remove first row
    refine_words_sheet = refine_words_sheet.iloc[1:]
    # column names 高情商-{target} -> target
    refine_words_sheet.columns = refine_words_sheet.columns.str.extract(r"高情商-(.*)")[
        0
    ]

    refine_words = {}
    for idx, col in enumerate(refine_words_sheet.columns):
        if not isinstance(col, str) or "Unnamed" in col:
            continue
        words = refine_words_sheet[col].dropna().tolist()
        types = refine_words_sheet.iloc[:, idx + 1].dropna().tolist()
        refine_words[col] = [
            RefineWord(word=word, type=types[i]) for i, word in enumerate(words)
        ]

    tag_sheet = xls.parse("tags")

    target = "恋人"
    tag_objs = {
        "恋人": [],
    }

    for _, row in tag_sheet.iterrows():
        # TODO spell error
        personality = row["Personility"]
        tone = row["Tone"]
        style = row["Style"]

        tag_objs[target].append(
            PreferenceTag(
                tag=row.iloc[0], tone=tone, style=style, personality=personality
            )
        )
    return tag_objs, refine_words
